import numpy as np


class Cartesian_coordinates_function:
    def __init__(self,space,uh):
        "实现将重新坐标变为笛卡尔坐标"
        mesh = space.mesh
        cell2dof = space.cell_to_dof()
        basis = space.basis

        cell = mesh.entity('cell') #(NC,gdim+1)
        node = mesh.entity('node') #(NN,gdim)
        dim = node.shape[-1]
        NC = mesh.number_of_cells()
        Node = node[cell] #(NC,gdim+1,gdim)
        #print(Node.shape)
        node0 = Node[:,0,:][:,None,:]
        matrix = Node[:,1:,:] - node0
        self.dim = dim
        self.mesh = mesh
        self.matrix = matrix
        self.node0 = node0
        self.NC = NC
        self.basis = basis
        self.uh = uh
        self.cell2dof = cell2dof
        self.space = space

    def point_to_index_and_bc(self,p): #p.shape = (Np,dim)
        dim = self.dim
        if p.shape[-1] != dim:
            print('dimension mismatch')
            return
        NC = self.NC
        matrix = self.matrix
        node0 = self.node0.T #(dim,1,NC)
        Np = p.shape[0] #要判断点的个数
        p = p.T #(dim,Np)
        idx = np.arange(Np) #未判断出来单元的点
        index = np.zeros(Np,dtype=int) #点对应点单元
        bc = np.zeros((Np,dim+1),dtype=float) #重心坐标


        for i in range(NC):
            A = matrix[i,:].T
            node = node0[...,i]
            #print(A)
            lam = np.linalg.solve(A,p[:,idx]-node) #(dim,...)
            if dim == 2:
                idxs = (((lam[0,:] > -1e-10)&(lam[1,:] > -1e-10))&((lam[0,:]+lam[1,:])<1+1e-10))
            elif dim == 3:
                idxs = (((lam[0,:] > -1e-10)&(lam[1,:] > -1e-10))
                        &((lam[2,:] > -1e-10)&((lam[0,:]+lam[1,:]+lam[2,:])<1+1e-10)))
            lam = lam[:,idxs].T
            bc[idx[idxs],1:] = lam
            bc[idx[idxs],0] = 1.0 - np.sum(lam,axis=-1)
            index[idx[idxs]] = i
            idx = np.delete(idx,idxs)
            if len(idx) == 0:
                break

        if len(idx) !=0:
            print('some points are wrang!')
            return
        return index, bc

    def value(self,p):
        dim = self.dim
        if p.shape[-1] != dim:
            print('dimension mismatch')
            return
        if len(p.shape) == 1:
            p = p[None,:] #只有一个点
        shape = p.shape[:-1]
        p = p.reshape((-1,dim)) #(Np.dim)
        Np = p.shape[0]
        index, bc = self.point_to_index_and_bc(p) #(Np,) (Np,dim+1)
        uh = self.uh
        cell2dof = self.cell2dof
        uh = uh[cell2dof[index]] #(Np,ldof)

        basis = self.basis(bc) 
        basis_shape = basis.shape #(Np,1,ldof,value_dim) or (Np,NC,ldof,value_dim)
        if basis_shape[1] == 1:
            val = np.einsum('ij,ij...->i...',uh,basis[:,0,...]) #(Np,...)
        else:
            val = np.einsum('ij,iij...->i...',uh,basis[:,index,...])
        shape += val.shape[1:]
        #print(shape)
        return val.reshape(shape)


    def values(self,p):
        dim = self.dim
        if p.shape[-1] != dim:
            print('dimension mismatch')
            return
        if len(p.shape) == 1:
            p = p[None,:] #只有一个点
        shape = p.shape[:-1]
        p = p.reshape((-1,dim)) #(Np.dim)
        Np = p.shape[0]
        index, bc = self.point_to_index_and_bc(p) #(Np,) (Np,dim+1)
        uh = self.uh
        cell2dof = self.cell2dof
        uh = uh[cell2dof[index]] #(Np,ldof,tdim)

        basis = self.basis(bc) #(Np,NC,ldof,tdim,value_dim)
        basis_shape = basis.shape #(Np,1,ldof,tdim,value_dim) or (Np,NC,ldof,tdim,value_dim)
        if basis_shape[1] == 1:
            val = np.einsum('ijk,ijk...->i...',uh,basis[:,0,...]) #(Np,...)
        else:
            val = np.einsum('ijk,iijk...->i...',uh,basis[:,index,...])
        shape += val.shape[1:]
        #print(shape)
        return val.reshape(shape)

    
    def grad_value(self,p):
        dim = self.dim
        if p.shape[-1] != dim:
            print('dimension mismatch')
            return
        if len(p.shape) == 1:
            p = p[None,:] #只有一个点
        shape = p.shape[:-1]
        p = p.reshape((-1,dim)) #(Np.dim)
        Np = p.shape[0]
        index, bc = self.point_to_index_and_bc(p) #(Np,) (Np,dim+1)
        uh = self.uh
        cell2dof = self.cell2dof
        uh = uh[cell2dof[index]] #(Np,ldof)

        grad_basis = self.space.grad_basis(bc) 
        grad_basis_shape = grad_basis.shape #(Np,NC,ldof,gdim,gdim)
        #print(grad_basis_shape)
        if grad_basis_shape[1] == 1:
            val = np.einsum('ij,ij...->i...',uh,grad_basis[:,0,...]) #(Np,...)
        else:
            val = np.einsum('ij,iij...->i...',uh,grad_basis[:,index,...])
        shape += val.shape[1:]
        #print(shape)
        return val.reshape(shape)

    def Recover_divu(self,p): #恢复div u
        grad_u_val_h = self.grad_value(p)
        return grad_u_val_h[...,0,0]+grad_u_val_h[...,1,1]
    def Recover_sgimah(self,p,lam=1,mu=0.5): #通过u恢复sigma
        grad_u_val_h = self.grad_value(p)
        div_u_val_h = grad_u_val_h[...,0,0]+grad_u_val_h[...,1,1]

        eplison_u_val_h = np.copy(grad_u_val_h)
        eplison_u_val_h[...,1,0] = 0.5*(eplison_u_val_h[...,1,0]+eplison_u_val_h[...,0,1])
        eplison_u_val_h[...,0,1] = eplison_u_val_h[...,1,0]

        sigma_h = 2*mu* eplison_u_val_h
        sigma_h[...,0,0]+=lam*div_u_val_h
        sigma_h[...,1,1]+=lam*div_u_val_h
        sigma_h = sigma_h[...,[0,1,1],[0,1,0]]
        
        return sigma_h
    
        
        
        
            



